// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/dwio/dwrf/reader/stripe_dictionary_cache.h>

namespace kumo::pollux::dwrf {
StripeDictionaryCache::DictionaryEntry::DictionaryEntry(
    melon::Function<BufferPtr(pollux::memory::MemoryPool*)>&& dictGen)
    : dictGen_{std::move(dictGen)} {}

BufferPtr StripeDictionaryCache::DictionaryEntry::getDictionaryBuffer(
    pollux::memory::MemoryPool* pool) {
  melon::call_once(onceFlag_, [&]() {
    dictionaryBuffer_ = dictGen_(pool);
    dictGen_ = nullptr;
  });
  return dictionaryBuffer_;
}

StripeDictionaryCache::StripeDictionaryCache(pollux::memory::MemoryPool* pool)
    : pool_{pool} {}

// It might be more elegant to pass in a StripeStream here instead.
void StripeDictionaryCache::registerIntDictionary(
    const EncodingKey& encodingKey,
    melon::Function<BufferPtr(pollux::memory::MemoryPool*)>&& dictGen) {
  intDictionaryFactories_.emplace(
      encodingKey, std::make_unique<DictionaryEntry>(std::move(dictGen)));
}

BufferPtr StripeDictionaryCache::getIntDictionary(
    const EncodingKey& encodingKey) {
  return intDictionaryFactories_.at(encodingKey)->getDictionaryBuffer(pool_);
}

} // namespace kumo::pollux::dwrf
